package wsz.gateway.filter;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 请求频率限流
 * @author wsz
 * @date 2021/12/29 9:36
 **/
@Component
@RefreshScope
public class RateLimitFilter implements GlobalFilter, Ordered {

    private static final String WARNING_MSG = "您频繁进行注册，请求已被拒绝";
    private static final Map<String, Bucket> LOCAL_CACHE = new ConcurrentHashMap<>();

    // 初始容量
    private long capacity = 1;

    @Value("${IP.minutes}")
    private long minutes;

    @Value("${IP.times}")
    private long times;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        ServerHttpResponse response = exchange.getResponse();
        String path = request.getURI().getPath();

        // TODO 注册-IP限流 使用验证码发送作为例子。
        if (path.contains("/user/register") || path.contains("/code/create")) {
            String ip = request.getRemoteAddress().getHostString();
            Bucket bucket = LOCAL_CACHE.computeIfAbsent(ip, k -> createNewBucket());
            if (!bucket.tryConsume(1)) {
                // 获取一个令牌失败
                response.setStatusCode(HttpStatus.FORBIDDEN);
                response.getHeaders().add("Content-Type", "text/plain;charset=UTF-8");
                return response.writeWith(Mono.just(response.bufferFactory().wrap(WARNING_MSG.getBytes(StandardCharsets.UTF_8))));
            }
        }

        return chain.filter(exchange);
    }

    private Bucket createNewBucket() {
        Refill refill = Refill.of(times, Duration.ofMinutes(minutes));
        Bandwidth limit = Bandwidth.classic(capacity, refill);
        return Bucket4j.builder().addLimit(limit).build();
    }

    @Override
    public int getOrder() {
        return 0;
    }
}
